
from Segmenter import Segmenter
from Model import TrackSegment  # noqa
from Model import Evidence  # noqa
import numpy as np
import copy
import sys
import os
import xml.etree.ElementTree as elet


class ChangepointFinder(Segmenter):
    """
    Reads in locations from a track as they become available, and looks for
    changepoints in the actor's movement regime given a set of parameters to be
    calculated from the track data and a set of models corresponding to known
    movement regimes (e.g. move fast vs move slow).

    Has as properties the functions for calculating each parameter given a set
    of track locations and the models for each pre-specified regime governing
    the actor's movement, as well as a variable indicating which model is the
    one of interest (i.e. for a "move fast" changepoint finder, we want to
    report the intervals where move fast is true)
    """

    current_frame = 1
    frames_to_param = dict()
    
    def __init__(self, **kwargs):
        kwarg_elts = ["param_func", "stat_funcs", "predictive_func",
                      "chi_prior", "nu_prior", "actor", "n_frames_for_params",
                      "hazard_rate", "false_chi", "false_nu", "true_chi",
                      "true_nu", "model_func", "model_params"]
        for e in kwarg_elts:
            if e in kwargs:
                setattr(self, e, kwargs[e])
        if "actor" in kwargs:
            initial_node = Node(run_length=1, time=0,
                                actor_start_time=self.actor.starttime)
            self.__current_messages = dict()
            self.__current_messages[str(initial_node)] = (
            initial_node, Message(
                weight=1, chi=self.chi_prior, nu=self.nu_prior, parent=None))
            self.__all_messages = copy.copy(self.__current_messages)
        self.__previous_frames = []
        self.__output_dir = "./changepoints"

    def print_changepoints(self):
        for changepoint in self.__changepoints:
            print changepoint

    def build_changepoint_list(self):
        best_path_weight = float("-inf")
        best_path_data = None
        for k, v in self.__current_messages.iteritems():
            node = v[0]
            message = v[1]
            # print "checking node %s with weight %s" % (node, message.weight)
            if message.weight > best_path_weight:
                # print "\n\nfound a better path!!\n\n"
                best_path_weight = message.weight
                best_path_data = (node, message)
        current_node = best_path_data[0]
        current_message = best_path_data[1]
        node_list = []
        while not current_node is None:
            current_message = self.__all_messages[str(
                    current_node)][1]
            node_list.append((current_node, current_message))
            current_node = Node(run_length=1, time=current_node.time -
                                (current_node.run_length - 1),
                                actor_start_time=self.actor.starttime)
            current_message = self.__all_messages[str(
                    current_node)][1]
            if not current_node is None:
                current_node = current_message.parent
        node_list.reverse()
        return node_list

    def fit_regimes(self):
        node_list = self.build_changepoint_list()
        regime_list = []
        for n, m in node_list:
            # want to fit regime to this segment
            chi_data = m.chi
            nu_data = m.nu
            end_time = n.time
            start_time = n.time - (n.run_length - 1)
            max_probability = float("-inf")
            max_label = ""
            # print ("\nprobabilities for interval [%s, %s] ([%s, %s]), with"
            #        "chi=%s, nu=%s, and chi/nu=%s" %
            #        (start_time, end_time, start_time+self.actor.starttime,
            #         end_time+self.actor.starttime,  chi_data[0], nu_data,
            #         chi_data[0]/nu_data))
            for p in self.model_params:
                param = p[1]
                label = p[0]
                cur_probability = self.model_func(chi_data, nu_data, param)
                # print ("probability %s for model with label %s" %
                #        (cur_probability, label))
                if cur_probability > max_probability:
                    max_probability = cur_probability
                    max_label = label
            regime_list.append(((start_time, end_time), max_label))
        return regime_list

    def join_regimes(self, regime_list):
        previous_regime = None
        result_list = []
        aggregated_interval = (-1, -1)
        for r in regime_list:
            start_time = r[0][0]
            end_time = r[0][1]
            regime = r[1]
            if previous_regime is not None:
                if regime == previous_regime:
                    aggregated_interval = (aggregated_interval[0], end_time)
                else:
                    result_list.append((aggregated_interval, previous_regime))
                    aggregated_interval = (start_time, end_time)
            else:
                aggregated_interval = (start_time, end_time)
            previous_regime = regime
        result_list.append((aggregated_interval, previous_regime))
        return result_list

    def generate_interval_xml(self, interval):
        start_time = interval[0][0] + self.actor.starttime
        end_time = interval[0][1] + self.actor.starttime
        label = interval[1]
        tag = "interval"
        adict = dict()
        adict["start_time"] = str(start_time)
        adict["end_time"] = str(end_time)
        adict["label"] = label
        elt = elet.Element(tag, adict)
        return elt

    def generate_changepoint_xml(self):
        regime_list = self.fit_regimes()
        regime_list = self.join_regimes(regime_list)
        adict = dict()
        adict["actor"] = self.actor.name
        wrapper_elt = elet.Element("intervals", adict)
        for r in regime_list:
            xml = self.generate_interval_xml(r)
            wrapper_elt.append(xml)
        return wrapper_elt
        
    def write_changepoints_to_file(self):
        def indent(elem, level=0):
            i = "\n" + level*"  "
            if len(elem):
                if not elem.text or not elem.text.strip():
                    elem.text = i + "  "
                if not elem.tail or not elem.tail.strip():
                    elem.tail = i
                for elem in elem:
                    indent(elem, level+1)
                if not elem.tail or not elem.tail.strip():
                    elem.tail = i
            else:
                if level and (not elem.tail or not elem.tail.strip()):
                    elem.tail = i
        wrapper_elt = self.generate_changepoint_xml()
        tree = elet.ElementTree(wrapper_elt)
        root = tree.getroot()
        indent(root)
        out_file = os.path.join(self.__output_dir, self.actor.name)
        # print "output file: %s" % out_file
        out_file = open(out_file, "w")
        tree.write(out_file)
        out_file.close()
            
    def output_current_changepoints(self):
        best_path_weight = float("-inf")
        best_path_data = None
        for k, v in self.__current_messages.iteritems():
            node = v[0]
            message = v[1]
            # print "checking node %s with weight %s" % (node, message.weight)
            if message.weight > best_path_weight:
                # print "\n\nfound a better path!!\n\n"
                best_path_weight = message.weight
                best_path_data = (node, message)
        current_node = best_path_data[0]
        current_message = best_path_data[1]
        count = 0
        while not current_node is None:
            # print "%s, %s" % (current_node, current_message)
            current_data = self.__all_messages[str(current_node)]
            current_message = current_data[1]
            if current_node.run_length == 0:
                p_string = ("\nat changepoint (time %s, frame %s) with data:")
                print ( p_string % (current_node.time, current_node.time +
                                    self.actor.starttime))
                print ("%s" % current_node)
                print ("%s" % current_message)
                current_node = current_message.parent
                if current_node is not None and current_node.time > 0:
                    print ("param: %s" %
                           ChangepointFinder.frames_to_param[
                               current_node.time])
                temp_node = current_node
                temp_message = current_message
                for i in range(5):
                    temp_node = temp_message.parent
                    if temp_node is not None:
                        temp_message = self.__all_messages[str(
                            temp_node)][1]
                        print temp_message
                        print ("param: %s" %
                               ChangepointFinder.frames_to_param[
                                   temp_node.time])
            else:
                p_string = ("\nat last node in regime (time %s, frame %s) with "
                            "data:")
                print ( p_string % (current_node.time, current_node.time +
                                    self.actor.starttime))
                print ("%s" % current_node)
                print ("%s" % current_message)
                current_node = Node(run_length=0, time=current_node.time -
                                    current_node.run_length,
                                    actor_start_time=self.actor.starttime)
                if current_node is not None and current_node.time > 0:
                    print ("param: %s" %
                           ChangepointFinder.frames_to_param[
                               current_node.time])
                temp_node = current_node
                temp_message = current_message
                for i in range(5):
                    temp_node = temp_message.parent
                    if temp_node is not None:
                        temp_message = self.__all_messages[str(
                            temp_node)][1]
                        print temp_message
                        print ("param: %s" %
                               ChangepointFinder.frames_to_param[
                                   temp_node.time])
            # print "found previous node %s" % prev_node
            

    def add_segment(self, segment):
        """
        Wrapper for __add_frames that handles evidence objects instead of a list
        of locations; exists so that all local inference classes can interface
        through evidence and actor objects without directly passing around
        location vectors

        "segment" is an evidence object; add_segment checks to make sure
        "segment"'s actor matches this CPF's actor, then extracts the
        "segment"'s list of frames and adds them to our list of frames for this
        actor and checks for changepoints
        """
        if segment.actor.name == self.actor.name and len(segment.locations) > 0:
            # don't add empty segment; actor check should always be true based
            # on subscriptions, but check anyway

            # print "adding segment of length %s" % len(segment.locations)
            # print "adding segment!"
            # print ("segment has the following location vector: %s"
            #        % segment.locations)
            self.__add_frames(segment.locations)

    def __add_frames(self, frame_list):
        """
        takes a list of locations, iterates over locations by frame and
        generates messages, then adds changepoints when the highest weight child
        node is a changepoint

        "frame_list" is the list of locations, indexed by frame offset from the
        first location
        """
        
        # print "in __add_frames with location vector %s" % frame_list

        if len(frame_list) < self.n_frames_for_params:
            # don't process frames if we don't have enough to calculate the
            # parameters, unless we have enough frames from before (otherwise
            # store these and wait for more)
            if (len(frame_list) + len(self.__previous_frames) <
                self.n_frames_for_params):
                if len(frame_list) > 0:
                    self.__previous_frames += frame_list
                return 0

        aggregate_frame_list = self.__previous_frames
        if len(aggregate_frame_list) < self.n_frames_for_params:
            # make sure aggregate_frame_list is long enough to calculate
            # parameters on first processed frame
            # print ("skipping %s frame(s) to set up parameter calculation" %
            #        (self.n_frames_for_params-1-len(aggregate_frame_list)))
            first_frame = (
                self.n_frames_for_params-1-len(aggregate_frame_list))
            # ChangepointFinder.current_frame += first_frame
            aggregate_frame_list += frame_list[:first_frame]
            frame_list = frame_list[first_frame:]

        messages = self.__current_messages
        n_frames = len(frame_list)
        for t, frame in enumerate(frame_list):
            # print ("***processing frame %s out of %s total***" %
            #         (t+1, n_frames))
            aggregate_frame_list += [frame]
            frame_param = self.__compute_param(aggregate_frame_list)
            ChangepointFinder.frames_to_param[
                ChangepointFinder.current_frame] = frame_param
            ChangepointFinder.current_frame += 1
            # print ("frame param is %s based on frames %s, %s" %
            #        (frame_param, aggregate_frame_list[-2],
            #         aggregate_frame_list[-1]))
            frame_stats = self.__calculate_statistics(frame_param)
            # print "frame stats are %s" % frame_stats
            predictive_prior = self.__compute_predictive_prior(frame_stats)
            new_messages = dict()
            n_nodes = len(messages)
            i = 0
            for k, v in messages.iteritems():
                node = v[0]
                message = v[1]
                # print ("***processing node %s out of %s total***"
                #        % (i+1, n_nodes))
                run_length = node.run_length
                weight, chi_prev, nu_prev = message.unpack()
                chi_growth = self.__update_statistics(chi_prev, frame_stats)
                predictive_posterior = self.__compute_predictive_posterior(
                    chi_prev, nu_prev, chi_growth)
                total_weight = (np.exp(predictive_posterior) +
                                     np.exp(predictive_prior))
                total_weight = np.log(total_weight)
                normed_post = predictive_posterior - total_weight
                normed_prior = predictive_prior - total_weight
                message_dict = self.__generate_child_messages(
                    node, self.hazard_rate, weight, normed_post,
                    normed_prior, frame_stats, chi_prev, nu_prev, node)
                self.__update_message_dict(new_messages, message_dict)
                i += 1
            messages = new_messages
            self.__update_message_dict(self.__all_messages, messages)
        self.__current_messages = messages
        self.__previous_frames = aggregate_frame_list

    def __compute_predictive_prior(self, statistics):
        updated_stats = self.__update_statistics(self.chi_prior, statistics)
        return self.predictive_func(self.chi_prior, self.nu_prior,
                                    updated_stats)

    def __compute_predictive_posterior(self, chi_prev, nu_prev, statistics):
        return self.predictive_func(chi_prev, nu_prev, statistics)

    def __compute_param(self, current_data):
        param_evidence = TrackSegment(
            actor=self.actor, segtype=None, starttime=0, endtime=0,
            locations=current_data)
        # param_evidence.locations = current_data
        # param_val = []
        # for p in self.param_funcs:
        #     param_vals.append(p(param_evidence))
        # return param_vals
        return self.param_func(param_evidence)

    def __update_statistics(self, current_stats, new_stats):
        # print ("updating statistics; current stats are %s; new stats are %s" %
        #        (current_stats, new_stats))
        updated_stats = []
        for i, s in enumerate(new_stats):
            updated_stats.append(current_stats[i] + s)
        return updated_stats
        
    def __calculate_statistics(self, parameter):
        stats = []
        for s in self.stat_funcs:
            chi = s(parameter)
            stats.append(chi)
        return stats

    def __update_message_dict(self, current_dict, new_dict):
        for k, v in new_dict.iteritems():
            if k in current_dict:
                message = v[1]
                best_message = current_dict[k][1]
                if message.weight > best_message.weight:
                    current_dict[k] = v
            else:
                current_dict[k] = v
        
    def __generate_child_messages(
            self, node, estimated_hazard_rate,
            prev_weight, predictive_post, predictive_prior, frame_stats,
            chi_prev, nu_prev, parent):
        """
        This function generates node objects corresponding to the children of
        the calling node and generates messages to be passed to each one.  Each
        node has two children: the growth child is the node at time t+1 when we
        do not predict a changepoint, while the change child is the child at
        time t+1 when we do predict a changepoint.  The growth child is unique
        to this node, while every node at time t shares a change child.  Since
        we are storing only the best path to each node, the change child will
        use only the highest scoring message it receives and pass that weight to
        its children.
        
        Return value is a dictionary in the same format as the node list
        so that its return value can easily be aggregated across nodes at a
        particular time step

        --the node's time and run_length are used to set the child's time and
        run_length values
        --n_changepoints, prev_weight, and nu_prev are properties of the calling
        node
        --estimated_hazard_rate is calculated based off of the calling node's
        n_changepoints value ((a_t+1)/(a_t+b_t+2) where a_t=n_c, b_t=t-a_t)
        --predictive_post is p(x_t|chi_parent, nu_parent)
        --predictive_prior is p(x_t|chi_0, nu_0)
        --chi_growth is chi_parent + T(x_t) (T is the set of sufficient
        statistics used to represent the distribution we're finding changepoints
        in)
        --parent is used for backtracking in order to output best path at a
        specified time
        --***estimated_hazard_rate is a constant for now, since including it
        makes the transition probabilities non-markovian***
        --***n_changepoints is currently not included since they're used to
        update the hazard rate estimation***
        """
        run_length = node.run_length
        c_time = node.time
        message_dict = dict()
        change_weight = (prev_weight + predictive_prior +
            np.log(estimated_hazard_rate))
        growth_weight = (prev_weight + predictive_post +
            np.log(1-estimated_hazard_rate))
        change_node = Node(run_length=1, time=c_time+1,
                           actor_start_time=self.actor.starttime)
        chi_change = self.__update_statistics(self.chi_prior, frame_stats)
        chi_growth = self.__update_statistics(chi_prev, frame_stats)
        change_message = Message(
            weight=change_weight, chi=chi_change, nu=self.nu_prior,
            parent=parent)
        message_dict[str(change_node)] = (change_node, change_message)
        growth_node = Node(run_length=run_length+1, time=c_time+1,
                           actor_start_time = self.actor.starttime)
        growth_message = Message(
            weight=growth_weight, chi=chi_growth, nu=nu_prev+1, parent=parent)
        message_dict[str(growth_node)] = (growth_node, growth_message)
        return message_dict

class Node(object):
    """
    Holds properties of a node in the message graph built by this algorithm

    Properties:
    run_length: the time elapsed since the last changepoint for this node
    n_changepoints: the number of changepoints observed overall for this
    node
    time: the current time
    """

    printed = dict()

    def __init__(self, **kwargs):
        kwarg_elts = ["run_length", "time", "actor_start_time"]
        for elt in kwarg_elts:
            if elt in kwargs:
                if elt == "time":
                    time = kwargs[elt]
                    if not time in Node.printed.keys():
                        # print "creating node with time %s" % kwargs[elt]
                        Node.printed[time] = 1
                setattr(self, elt, kwargs[elt])

    def __str__(self):
        return ("**Node:[run_length=%s,time=%s(%s)]**" %
                (self.run_length, self.time, self.time+self.actor_start_time))

class Message(object):
    """
    """

    def __init__(self, **kwargs):
        kwarg_elts = ["n_changepoints", "weight", "chi", "nu", "parent"]
        for e in kwarg_elts:
            if e in kwargs:
                setattr(self, e, kwargs[e])

    def unpack(self):
        return self.weight, self.chi, self.nu

    def __str__(self):
        # print self.chi
        chi_string = "[%s]" % ",".join([str(x) for x in self.chi])
        return ("**Message:[weight=%s,chi=%s,nu=%s,parent=%s]**" %
                (self.weight, chi_string, self.nu, self.parent))
